# cron: 15 6 * * *
# new Env("51代理更换白名单")
import requests
import os
import sys
import time
from datetime import datetime
import warnings
from urllib3.exceptions import InsecureRequestWarning

# 禁用SSL警告
warnings.filterwarnings("ignore", category=InsecureRequestWarning)

# ================= 配置区 =================
API_KEY = os.getenv('API_KEY_51')  # 从环境变量获取
if not API_KEY:
    print("错误：未找到API_KEY_51环境变量")
    sys.exit(1)

# 从环境变量获取需要锁定的IP，多个IP用#号分割
FIXED_IPS_STR = os.getenv('FIXED_IPS_51', '').strip()
if FIXED_IPS_STR:
    FIXED_IPS = set(ip.strip() for ip in FIXED_IPS_STR.split('#') if ip.strip())
else:
    FIXED_IPS = set()

SUCCESS_CODES = [0, 1, 200]    # 接口成功状态码
MAX_RETRIES = 3               # 最大重试次数
BASE_DELAY = 2                # 基础延迟(秒)
# ==========================================

def debug_log(msg, level="INFO"):
    """分级日志系统"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    color_code = {
        "DEBUG": "\033[94m",   # 蓝色
        "INFO": "\033[92m",    # 绿色
        "WARNING": "\033[93m", # 黄色
        "ERROR": "\033[91m",   # 红色
        "SUCCESS": "\033[96m", # 青色
        "SYSTEM": "\033[95m"   # 紫色
    }.get(level, "\033[0m")
    print(f"{color_code}## [{timestamp}] [{level}] {msg}\033[0m")

def get_public_ip():
    """实时获取公网IP（禁用缓存）"""
    services = [
        {'url': 'https://api.ipify.org', 'verify': True, 'timeout': 3},
        {'url': 'https://ident.me', 'verify': False, 'timeout': 3},
        {'url': 'https://ipinfo.io/ip', 'verify': True, 'timeout': 3},
        {'url': 'http://myip.ipip.net', 'verify': False, 'timeout': 3}
    ]
    
    for service in services:
        try:
            response = requests.get(
                service['url'],
                timeout=service['timeout'],
                verify=service['verify']
            )
            if response.status_code == 200:
                ip = response.text.strip()
                if validate_ip(ip):
                    debug_log(f"从 {service['url']} 获取到有效公网IP：{ip}", "SUCCESS")
                    return ip
        except Exception as e:
            debug_log(f"{service['url']}查询失败：{str(e)}", "DEBUG")
            continue
            
    debug_log("所有IP服务不可用", "ERROR")
    return None

def validate_ip(ip):
    """验证IP格式有效性"""
    parts = ip.split('.')
    if len(parts) != 4:
        return False
    try:
        return all(0 <= int(part) < 256 for part in parts)
    except ValueError:
        return False

def get_whitelist_with_details():
    """获取当前白名单列表（带详细信息）"""
    debug_log("正在获取白名单详情...", "DEBUG")
    time.sleep(BASE_DELAY)
    
    url = 'http://bapi.51daili.com/white-ip/list'
    try:
        response = requests.get(url, params={'appkey': API_KEY}, timeout=10)
        debug_log(f"接口响应状态码：{response.status_code}", "DEBUG")
        
        data = response.json()
        
        # 检查响应结构
        if 'code' not in data:
            debug_log("API响应缺少code字段", "ERROR")
            return None
            
        if data.get('code') in SUCCESS_CODES:
            # 检查rows字段是否存在
            if 'rows' not in data:
                debug_log("API响应缺少rows字段", "ERROR")
                return None
                
            ip_details = []
            for item in data.get("rows", []):
                if isinstance(item, dict) and "whiteip" in item and item["whiteip"]:
                    ip = item["whiteip"].strip()
                    if validate_ip(ip):
                        # 获取更新时间（根据API响应，字段名是updatetime）
                        update_time = item.get("updatetime", "")
                        ip_details.append({
                            "ip": ip,
                            "update_time": update_time
                        })
            
            debug_log(f"成功获取到{len(ip_details)}个有效白名单IP", "SUCCESS")
            return ip_details
        else:
            error_msg = data.get('msg', '未知错误')
            debug_log(f"API返回错误：{error_msg} (代码: {data.get('code')})", "ERROR")
            return None
            
    except Exception as e:
        debug_log(f"请求白名单接口失败：{str(e)}", "ERROR")
        return None

def get_whitelist():
    """获取当前白名单IP集合"""
    ip_details = get_whitelist_with_details()
    if ip_details:
        return set(item["ip"] for item in ip_details)
    return None

def delete_one_oldest_non_fixed_ip(current_ips, public_ip):
    """删除最久的一个非锁定IP"""
    # 获取详细的白名单信息
    ip_details = get_whitelist_with_details()
    if not ip_details:
        return False
    
    # 找出所有非锁定IP（不包括当前公网IP）
    non_fixed_ips = []
    for item in ip_details:
        ip = item["ip"]
        if ip != public_ip and ip not in FIXED_IPS:
            non_fixed_ips.append(item)
    
    if not non_fixed_ips:
        debug_log("没有找到可删除的非锁定IP", "WARNING")
        return False
    
    # 如果有更新时间信息，按时间排序；否则任意选择一个
    if all("update_time" in item and item["update_time"] for item in non_fixed_ips):
        # 按更新时间排序（最早的优先）
        non_fixed_ips.sort(key=lambda x: x["update_time"])
        ip_to_delete = non_fixed_ips[0]["ip"]
        debug_log(f"找到最久的非锁定IP: {ip_to_delete} (更新时间: {non_fixed_ips[0]['update_time']})", "INFO")
    else:
        # 没有时间信息，任意选择一个
        ip_to_delete = non_fixed_ips[0]["ip"]
        debug_log(f"选择删除非锁定IP: {ip_to_delete}", "INFO")
    
    # 执行删除
    return delete_ips([ip_to_delete])

def delete_ips(ips):
    """删除指定IP列表"""
    if not ips:
        debug_log("无需删除IP", "INFO")
        return True
    
    debug_log(f"准备删除{len(ips)}个IP: {', '.join(ips)}", "INFO")
    time.sleep(BASE_DELAY)
    
    url = 'http://bapi.51daili.com/white-ip/del'
    try:
        params = {
            'appkey': API_KEY,
            'ips': ','.join(ips)
        }
        
        response = requests.get(url, params=params, timeout=15)
        data = response.json()
        
        if data.get('code') in SUCCESS_CODES:
            debug_log(f"成功删除IP: {', '.join(ips)}", "SUCCESS")
            return True
        else:
            debug_log(f"删除失败：{data.get('msg', '未知错误')}", "ERROR")
            return False
    except Exception as e:
        debug_log(f"删除操作异常：{str(e)}", "ERROR")
        return False

def add_ip(ip):
    """添加IP到白名单"""
    debug_log(f"开始添加IP：{ip}", "INFO")
    time.sleep(BASE_DELAY)
    
    url = 'http://bapi.51daili.com/white-ip/add'
    params = {
        'appkey': API_KEY,
        'isdel': 1,
        'ips': ip
    }
    
    try:
        response = requests.get(url, params=params, timeout=15)
        data = response.json()
        
        if data.get('code') in SUCCESS_CODES:
            msg = data.get('msg', '')
            if "已存在" in msg:
                debug_log(f"IP {ip} 已存在", "INFO")
                return True
            debug_log(f"IP {ip} 添加成功", "SUCCESS")
            return True
            
        debug_log(f"添加失败：{data.get('msg', '未知错误')}", "ERROR")
        return False
    except Exception as e:
        debug_log(f"添加请求异常：{str(e)}", "ERROR")
        return False

def verify_ip(ip, max_retries=3):
    """严格验证IP是否存在"""
    debug_log(f"开始验证IP {ip} 是否生效", "INFO")
    for retry in range(max_retries):
        current_ips = get_whitelist()
        if not current_ips:
            debug_log(f"第{retry+1}次验证失败：无法获取白名单", "WARNING")
            time.sleep(BASE_DELAY * 2)
            continue
            
        if ip in current_ips:
            debug_log(f"第{retry+1}次验证通过", "SUCCESS")
            return True
            
        debug_log(f"第{retry+1}次验证未找到IP", "WARNING")
        time.sleep(BASE_DELAY * 2)
        
    debug_log("超过最大验证次数仍未生效", "ERROR")
    return False

def main():
    debug_log("===== 白名单管理流程开始 =====", "SYSTEM")
    
    # 检查环境变量
    debug_log(f"使用API密钥：{API_KEY[:8]}...", "DEBUG")
    
    # 显示锁定IP配置
    if FIXED_IPS:
        debug_log(f"配置的锁定IP：{', '.join(FIXED_IPS)}", "INFO")
    else:
        debug_log("未配置锁定IP，将只管理当前公网IP", "INFO")
    
    # 第一步：实时获取公网IP
    debug_log("正在获取公网IP...", "INFO")
    public_ip = get_public_ip()
    if not public_ip:
        debug_log("无法获取有效公网IP", "FATAL")
        sys.exit(1)
    debug_log(f"当前公网IP：{public_ip}", "SUCCESS")
    
    # 第二步：获取当前白名单（添加重试机制）
    current_whitelist = None
    for retry in range(MAX_RETRIES):
        current_whitelist = get_whitelist()
        if current_whitelist is not None:
            break
        if retry < MAX_RETRIES - 1:
            debug_log(f"第{retry+1}次获取白名单失败，等待重试...", "WARNING")
            time.sleep(BASE_DELAY * (retry + 1))
    
    if not current_whitelist:
        debug_log("无法获取白名单列表", "FATAL")
        sys.exit(1)
    
    debug_log(f"当前白名单：{current_whitelist}", "INFO")
    
    # 判断是否需要处理
    if public_ip in current_whitelist:
        debug_log("当前IP已在白名单中", "SUCCESS")
        
        # 如果有锁定IP，检查是否需要添加缺失的锁定IP
        if FIXED_IPS:
            missing_fixed_ips = FIXED_IPS - current_whitelist
            if missing_fixed_ips:
                debug_log(f"发现{len(missing_fixed_ips)}个缺失的锁定IP", "WARNING")
                for ip in missing_fixed_ips:
                    if not add_ip(ip):
                        debug_log(f"添加锁定IP {ip} 失败", "ERROR")
            else:
                debug_log("所有锁定IP已在白名单中", "SUCCESS")
        sys.exit(0)
        
    debug_log("当前IP不在白名单中，开始处理流程...", "WARNING")
    
    # 第三步：尝试添加新IP，如果失败则删除一个最久的非锁定IP
    if not add_ip(public_ip):
        debug_log("首次添加IP失败，尝试删除一个最久的非锁定IP后重试", "WARNING")
        if not delete_one_oldest_non_fixed_ip(current_whitelist, public_ip):
            debug_log("删除非锁定IP失败", "ERROR")
            sys.exit(1)
        
        # 再次尝试添加
        if not add_ip(public_ip):
            debug_log("再次添加IP失败", "FATAL")
            sys.exit(1)
    
    # 第四步：如果有锁定IP，确保它们都在白名单中
    if FIXED_IPS:
        missing_fixed_ips = FIXED_IPS - current_whitelist
        if missing_fixed_ips:
            debug_log(f"需要添加{len(missing_fixed_ips)}个缺失的锁定IP", "INFO")
            for ip in missing_fixed_ips:
                if not add_ip(ip):
                    debug_log(f"添加锁定IP {ip} 失败", "ERROR")
    
    # 第五步：严格验证当前公网IP
    if not verify_ip(public_ip):
        debug_log("最终验证失败", "FATAL")
        sys.exit(1)
    
    # 最终检查
    final_list = get_whitelist()
    if final_list:
        # 期望的白名单：当前公网IP + 锁定IP（如果有）
        expected_ips = {public_ip}
        if FIXED_IPS:
            expected_ips = expected_ips.union(FIXED_IPS)
        
        if final_list == expected_ips:
            debug_log("所有IP验证通过", "SUCCESS")
            debug_log(f"最终白名单：{', '.join(sorted(final_list))}", "INFO")
        else:
            debug_log(f"最终检查未通过，期望：{expected_ips}，实际：{final_list}", "ERROR")
            sys.exit(1)
    else:
        debug_log("无法获取最终白名单进行验证", "ERROR")
        sys.exit(1)
        
    debug_log("===== 操作成功完成 =====", "SYSTEM")

if __name__ == '__main__':
    main()